import concurrent.futures
import datetime
import json
import logging
import os
import tempfile
import traceback
import typing

import pandas as pd
from google.protobuf.json_format import MessageToDict

import coin.base.database.util as db_util
import coin.proto.coin_data_replay_pb2 as cdr_pb2
from coin.base.timestamp import get_timestamp
from coin.base.datetime_util import (to_datetime, to_timestamp_int)
from coin.base.param_util import to_list
from coin.proto.coin_market_query_pb2 import ProductTicker, ExchangeTicker
from coin.proto.coin_request_pb2 import (
    ExchangeApiInfoProto,
    AccountInfoProto,
    StrategyInfoProto,
    StrategyRequestProto)
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_strategy_summary_pb2 import (StrategyDetailEntry,
                                                  StrategyDetailSummary,
                                                  StrategyStat,
                                                  StrategyStatLog,
                                                  StrategySummaryMessage)
from coin.proto.coin_telemetry_pb2 import StrategyLatencyStatsProto
from coin.support.data_replay.driver import DataReplayHandler
from coin.support.proto_log.logic.util import StratInfo
from coin.support.pta.logging.strategy_stat_logger import (
    gen_strategy_stat_dir,
    StrategyStatReader,)
from coin.support.pta.logic.pta_adjustment import adjust_strat_log
from coin.support.pta.util.dump_util import dump_pta_to_csv
from coin.tool.strat_monitor.handlers.strategy_summary_calculator import \
    StrategySummaryCalculator
from coin.tool.strat_monitor.plot.strat_plotter import DailyStrategyPlotterConfig
from coin.tool.strat_monitor.util.display_util import (
    print_strategy_summary, send_strat_summary_to_slack)
from xunkemgmt_client.base.param_util import split_list_by_group
from xunkemgmt_client.support.pta.util.proto_util import convert_pta_proto_to_df
from xunkemgmt_client.tool.slack_noti import send_to_slack
from xunkemgmt_client.client.util.query_util import (
    query_accounts,
    query_exchange_apis,
    query_strategies)
from xunkemgmt_client.support.pta.database.mysql_importer import PtaStatsImporter


class _Onlog:
  def __init__(self):
    self.strat_stat_logs = []

  def on_log(self, timestamp: int, log):
    pb = StrategyStatLog()
    pb.ParseFromString(log)
    self.strat_stat_logs.append(pb)


class PtaDumpStrategy(DataReplayHandler):
  def __init__(self,
               *,
               strat_info: StratInfo,
               start_time: datetime.datetime,
               end_time: datetime.datetime,
               source_info: str=None,
               ignore_error: bool=False,
               exchange_api_universe: \
                   typing.Optional[typing.List[ExchangeApiInfoProto]]=None,
               account_universe: typing.Optional[typing.List[AccountInfoProto]]=None,
               strategy_universe: typing.Optional[typing.List[StrategyInfoProto]]=None,
               pnl_adj=None,
               realtime_pnl_adj: bool=False,
               ticker: typing.Optional[typing.Dict[str, float]]=None,
               plot_dir: typing.Optional[str]=None,
               plot_order: bool=False,
               save_order_event: bool=False,
               log_history_dir: typing.Optional[str]=None,
               strat_stat_log_dir: typing.Optional[str]=None):
    self._strat_info = strat_info
    self._strategy_name = strat_info.strategy_name
    self._machine = strat_info.machine
    self._start_ts = to_timestamp_int(start_time)
    self._end_ts = to_timestamp_int(end_time)
    self._curr_running_ts = None
    if plot_dir is not None:
      if strat_stat_log_dir is None:
        strat_stat_log_dir = tempfile.mkdtemp(prefix='pnl_stat_proto_log', dir='/tmp')
    self._pta_stats = StrategySummaryCalculator(strategy_name=self._strategy_name,
                                                machine=self._machine,
                                                source_info=source_info,
                                                start_timestamp=self._start_ts,
                                                current_timestamp=self._start_ts,
                                                daily_reset_disabled=True,
                                                save_order_event=save_order_event,
                                                save_history=True,
                                                strat_stat_log_dir=strat_stat_log_dir)
    self._exchange_api_universe = \
        query_exchange_apis(as_proto=True) if exchange_api_universe is None else \
            exchange_api_universe
    self._account_universe = \
        query_accounts(as_proto=True) if account_universe is None else account_universe
    self._strategy_universe = \
        query_strategies(as_proto=True) if strategy_universe is None else \
        strategy_universe
    self._ignore_error = ignore_error
    self._pnl_adj = pnl_adj
    self._realtime_pnl_adj = realtime_pnl_adj
    self._pta_stats.update_exchange_api_universe(self._exchange_api_universe)
    self._pta_stats.update_account_universe(self._account_universe)
    self._pta_stats.update_strategy_universe(self._strategy_universe)
    self._ticker = ticker if ticker is not None else {}
    # price must be available if calculator needs to adjust pnl by reserve diff instead
    # of pnl balance diff on reserve change. ticker in strat log is not available until
    # 20210924, to adjust pnl by reserve diff before this time, eod price is provided to
    # achieve this.
    if self._ticker:
      ticker_log = StrategyLog(
          type=StrategyLog.TICKER,
          ticker=ExchangeTicker(each_ticker=[ProductTicker(
              symbol=symbol, last=price) for symbol, price in self._ticker.items()]))
      self._pta_stats.handle(self._start_ts, ticker_log)
    self._plot_dir = plot_dir
    self._plot_order = plot_order
    self._log_history_dir = log_history_dir
    self._strat_stat_log_dir = strat_stat_log_dir

    self._pta_summaries = []
    self._pta_details = []
    self._logger = logging.getLogger(__name__)

  @property
  def strat_info(self):
    return self._strat_info

  @property
  def pta_summaries(self):
    return self._pta_summaries

  @property
  def pta_details(self):
    return self._pta_details

  def _do_pnl_plot(self):
    from coin.support.pta.plot.strat_plotter import StrategyPlotter
    acct_df = pd.DataFrame([MessageToDict(acct, preserving_proto_field_name=True)
                            for acct in self._account_universe])
    pnl_adj = self._pnl_adj if self._realtime_pnl_adj else None
    try:
      onlog = _Onlog()
      stat_reader = StrategyStatReader()
      start_dt = to_datetime(self._start_ts)
      end_dt = to_datetime(self._end_ts)
      stat_reader.run_from_archive(
          on_log_callback=onlog.on_log,
          start_time=start_dt,
          end_time=end_dt,
          root_dir=self._strat_stat_log_dir,
          machine=self._machine,
          strategy_name=self._strategy_name,
      )
      plot_dir = gen_strategy_stat_dir(
          log_root=self._plot_dir,
          trading_date=start_dt,
          strat_request=StrategyRequestProto(strategy_name=self._strategy_name),
          machine=self._machine)
      plotter = StrategyPlotter(acct_df, pnl_adj)
      plotter.do_pnl_plot(
          strat_name=self._strategy_name,
          strat_stat_logs=onlog.strat_stat_logs,
          start_time=start_dt,
          end_time=end_dt,
          plot_dir=plot_dir)
    except Exception:
      logging.error('Fail to dump pnl plot. %s, %s' % (
          self._strategy_name, self._machine))
      logging.error(traceback.format_exc())

  def _do_order_plot(self):
    from coin.tool.strat_monitor.plot.strat_plotter import StrategyPlotter
    start_date = to_datetime(self._start_ts).date()
    end_date = to_datetime(self._end_ts).date()
    plotter_config = DailyStrategyPlotterConfig(start_date, end_date, self._plot_dir)
    plotter = StrategyPlotter(plotter_config, account_info=self._account_universe)
    try:
      plotter.do_order_plot(to_list(self._pta_stats.get_order_plot_info()))
    except Exception:
      logging.error('Fail to dump order plot. %s, %s' % (
          self._strategy_name, self._machine))
      logging.error(traceback.format_exc())

  def on_replay_start(self, ts: int):
    self._curr_running_ts = ts

  def on_replay_end(self, ts: int):
    start_date = to_datetime(self._start_ts).date()
    end_date = to_datetime(self._end_ts).date()
    if self._plot_dir is not None:
      self._do_pnl_plot()
      if self._plot_order:
        self._do_order_plot()
    if self._log_history_dir is not None:
      for (account_request, symbol), log_history in \
          self._pta_stats.get_log_history().items():
        output_dir = get_order_event_data_dir(
            root_dir=self._log_history_dir,
            start_date=start_date,
            end_date=end_date,
            strategy_name=self._strategy_name,
            machine=self._machine,
            account_request=account_request,
            symbol=symbol)
        os.makedirs(output_dir, exist_ok=True)
        filename = 'order_event.json'
        with open('%s/%s' % (output_dir, filename), mode='w') as fd:
          json.dump(log_history, fd)


  def on_strat_log(self,
                   topic_info: cdr_pb2.TopicDataInfo,
                   ts: int,
                   pb: StrategyLog):
    try:
      adjust_strat_log(self._strategy_name, ts, pb)
      self._pta_stats.handle(ts, pb)
    except Exception as e:
      self._logger.info('')
      self._logger.info('%s, %s' % (ts, pb))
      if self._ignore_error:
        self._logger.error(traceback.format_exc())
      else:
        raise e

  def on_timestamp(self, ts: int):
    if ts > self._start_ts:
      # TODO(chensili): change it to use latest ticker before ts
      self._pta_stats.update_ticker(self._ticker, self._end_ts)
      pta_summary = self._pta_stats.get_result(ts)
      curr_running_ts = get_timestamp()
      duration = curr_running_ts - self._curr_running_ts
      pta_summary.source.strategy_log_info.log_replay_duration = duration
      self._pta_summaries.append(pta_summary)
      self._pta_details.append(self._pta_stats.get_history())


def get_order_event_data_dir(
    *, root_dir, start_date, end_date, strategy_name, machine, account_request, symbol):
  config = DailyStrategyPlotterConfig(start_date, end_date, root_dir)
  output_dir = config.get_order_plot_dir(
          strategy_name=strategy_name,
          machine=machine,
          account_request=account_request,
          symbol=symbol)
  output_dir = os.path.join(output_dir, 'data')
  return output_dir


def _insert_latency_stats(latency_stats: typing.List[StrategyLatencyStatsProto],
                          db_config_path: str):
  from xunkemgmt_client.support.telemetry.database.clickhouse.db_importer import \
      LatencyStatsImporter
  db_config = db_util.read_db_config(db_config_path)
  importer = LatencyStatsImporter.create_connector(db_config)
  importer.bulk_insert_strategy_latency_stats(latency_stats)


def dump_pta_stats(summary_msg,
                   strategy_details,
                   *,
                   ticker=None,
                   output_file=None,
                   output_csv_file=None,
                   slack_receiver=None,
                   pta_db_config_path=None,
                   latency_db_config_path=None):
  assert isinstance(summary_msg, StrategySummaryMessage)
  summary_string = print_strategy_summary(summary_msg, mode='report', ticker=ticker)
  if output_file is not None:
    with open(output_file, "w") as fp:
      fp.write(summary_string)
  send_strat_summary_to_slack(slack_receiver, summary_string)
  if pta_db_config_path is not None:
    db_config = db_util.read_db_config(pta_db_config_path)
    importer = PtaStatsImporter(db_config=db_config)
    for strategy_pta in summary_msg.summaries:
      if len(strategy_pta.stat_per_account) == 0 and \
         len(strategy_pta.stat_per_quote) == 0:
        continue
      importer.insert_strategy_pta_stats(strategy_pta)
    for strategy_detail in strategy_details:
      if len(strategy_detail.each_detail) > 0:
        importer.insert_strategy_detail(strategy_detail)
  if latency_db_config_path is not None:
    latency_stats = [summary.latency_stats_1min for summary in summary_msg.summaries]
    max_workers = 48
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
      futures = []
      for latency_stats in split_list_by_group(latency_stats, max_workers):
        future = executor.submit(_insert_latency_stats, latency_stats, latency_db_config_path)
        futures.append(future)
      for future in futures:
        future.result()
  if output_csv_file is not None:
    pta_df = convert_pta_proto_to_df(summary_msg)
    pta_df = pta_df[pd.isnull(pta_df['symbol'])]
    dump_pta_to_csv(pta_df, output_csv_file)
    with open(output_csv_file) as f:
      send_to_slack(f.read(), slack_receiver, 'file')


def aggregate_strategy_detail(detail, period_minute):
  assert isinstance(detail, StrategyDetailSummary)
  assert isinstance(period_minute, int)
  assert period_minute in (1, 5, 10, 15, 20, 30, 60), period_minute
  period = '1h' if period_minute == 60 else '%m' % period_minute
  agg_each_detail = []
  for each_detail in detail.each_detail:
    pnl_detail = []
    for each_pnl in each_detail.pnl:
      assert each_pnl.HasField('stat_ts'), each_pnl
      pnl_mark_from_balance = each_pnl.pnl_mark if each_pnl.HasField('pnl_mark_from_balance') else None
      pnl_mark_from_fill = each_pnl.pnl_mark_from_fill if \
                           each_pnl.HasField('pnl_mark_from_fill') else None
      pnl_detail.append([each_pnl.stat_ts, pnl_mark_from_balance, pnl_mark_from_fill])
    pnl_df = pd.DataFrame(pnl_detail, columns=['stat_ts', 'pnl_mark_from_balance', 'pnl_mark_from_fill'])
    pnl_df['stat_ts'] = pd.to_datetime(pnl_df['stat_ts'])
    pnl_df.set_index('stat_ts', inplace=True)
    pnl_df = pnl_df.resample('%dT' % period_minute).last()
    pnl_df.reset_index(inplace=True)
    agg_detail_entry = StrategyDetailEntry()
    agg_detail_entry.CopyFrom(each_detail)
    agg_detail_entry.ClearField('pnl')
    agg_detail_entry.ClearField('position')
    agg_pnl_detail = []
    for _, row in pnl_df.iterrows():
      pnl_mark_from_fill = row['pnl_mark_from_fill'] if not \
                           pd.isnull(row['pnl_mark_from_fill']) else None
      pnl_mark_from_balance = row['pnl_mark_from_balance'] if not \
                              pd.isnull(row['pnl_mark_from_balance']) else None
      agg_pnl_detail.append(
          StrategyStat(stat_ts=int(row['stat_ts'].timestamp() * 10**9),
                       pnl_mark_from_balance=pnl_mark_from_balance,
                       pnl_mark_from_fill=pnl_mark_from_fill))
    agg_detail_entry.pnl.extend(agg_pnl_detail)
    agg_each_detail.append(agg_detail_entry)
  agg_detail = StrategyDetailSummary(process=detail.process,
                                     strategy=detail.strategy,
                                     calculator=detail.calculator,
                                     source_info=detail.source_info,
                                     detail_period=period,
                                     each_detail=agg_each_detail)
  return agg_detail
